/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2007  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

//#define	TRACE_LEVEL	2


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#define _NTUNDOC_C
#include "NtUndoc.h"

// project's headers
#include "ImgInfo.h"
#include "ProcInfo.h"
#include "SymbolFile.h"
#include "SystInfo.h"
#include "Trace.h"


/******************************************************************/
/* Exported data                                                  */
/******************************************************************/

NTUNDOC_NAMESPACE ntundoc ;


/******************************************************************/
/* Internal macros                                                */
/******************************************************************/

#define arraysize(a) (sizeof(a)/sizeof((a)[0]))

#define SET_NTUNDOC_PROC(proc)						\
  {									\
    INT_PTR nOffset = SymFile_GetProcAddress(hSymFile,hSymKern,TEXT(#proc)) ; \
    if( nOffset != 0 )							\
      {									\
	ntundoc.proc = (VOID*)(pKernelBase+nOffset) ;			\
      }									\
    else								\
      {									\
	ntundoc.proc = NULL ;						\
	TRACE_WARNING(TEXT(#proc " address is NULL\n")) ;			\
      }									\
  }

#define COMPILE_TIME_ASSERT(pred)       \
    switch(0){case 0:case pred:;}

/******************************************************************/
/* Exported function                                              */
/******************************************************************/

NTSTATUS NtUndoc_Init() 
{
  NTSTATUS		nStatus ;
  IMGINFO		imgInfo ;
  CHAR			szFilename[MAX_PATH] ;
  PBYTE			pKernelBase = NULL ;
  ULONG			nKernelSize = 0 ;
  HSYMFILE		hSymFile ;
  HSYMKERN		hSymKern ;
  
  SystInfo_GetModuleBase (NULL, (void**)&pKernelBase, &nKernelSize) ;

  TRACE_INFO(TEXT("Kernel base : 0x%08X\n"), pKernelBase) ;
  TRACE_INFO(TEXT("Kernel size : 0x%08X\n"), nKernelSize) ;

  if( ! ImgInfo_GetInfo (&imgInfo, pKernelBase, nKernelSize, FALSE) )
    {
      TRACE_ERROR (TEXT("ImgInfo_GetInfo failed\n")) ;
      return STATUS_UNSUCCESSFUL ;
    }

  TRACE_INFO(TEXT("Image signature = %s\n"), imgInfo.szSymSignature) ;

  nStatus = SystInfo_GetThisImageName (szFilename) ;
  //nStatus = ProcInfo_GetCurrentProcessImagePath (&usFilename) ;
  if( nStatus != STATUS_SUCCESS ) return nStatus ;
 
  TRACE_INFO(TEXT("Driver path = %hs\n"), szFilename) ;

  if( strrchr(szFilename,'\\') == NULL )
    {
      TRACE_ERROR (TEXT("File name format invalid\n")) ;
      return STATUS_UNSUCCESSFUL ;
    }

  strcpy (strrchr(szFilename,'\\'), "\\Symbols.wps") ;
  
  TRACE_INFO(TEXT("Symbol database = %hs\n"), szFilename) ;

  hSymFile = SymFile_Load (szFilename) ;

  if( hSymFile == NULL )
    {
      TRACE_ERROR (TEXT("SymFile_Load failed\n")) ;
      return STATUS_UNSUCCESSFUL ;
    }

  hSymKern = SymFile_GetKernel (hSymFile, imgInfo.szSymSignature) ;

  if( hSymKern == NULL )
    {
      TRACE_ERROR (TEXT("SymFile_GetKernel failed\n")) ;
      SymFile_CloseFile (hSymFile) ;
      return STATUS_UNSUCCESSFUL ;
    }

  SET_NTUNDOC_PROC(ExfAcquirePushLockExclusive) ;
  SET_NTUNDOC_PROC(ExfReleasePushLock) ;
  SET_NTUNDOC_PROC(NtCreateProcess) ;
  SET_NTUNDOC_PROC(NtCreateProcessEx) ;
  SET_NTUNDOC_PROC(NtCreateSection) ;
  SET_NTUNDOC_PROC(NtCreateUserProcess) ;
  SET_NTUNDOC_PROC(NtTerminateProcess) ;
  SET_NTUNDOC_PROC(NtQueryInformationFile) ;
  SET_NTUNDOC_PROC(NtQueryKey) ;
  SET_NTUNDOC_PROC(NtQueryValueKey) ;
  SET_NTUNDOC_PROC(NtSetInformationFile) ;
  SET_NTUNDOC_PROC(NtSetValueKey) ;
  SET_NTUNDOC_PROC(ObpFreeObject) ;
  SET_NTUNDOC_PROC(PspTerminateProcess) ;
  SET_NTUNDOC_PROC(swprintf) ;
  SET_NTUNDOC_PROC(ZwOpenProcess) ;
  SET_NTUNDOC_PROC(ZwProtectVirtualMemory) ;
  SET_NTUNDOC_PROC(ZwReadVirtualMemory) ;

  COMPILE_TIME_ASSERT(sizeof(OBJECT_HEADER)==0x020) ;
  COMPILE_TIME_ASSERT(sizeof(OBJECT_TYPE)==0x0B0) ;
  COMPILE_TIME_ASSERT(sizeof(FILE_OBJECT_W2K)==0x70) ;

  COMPILE_TIME_ASSERT(sizeof(ALPC_PROCESS_CONTEXT)==0x10) ;

  COMPILE_TIME_ASSERT(sizeof(MMSUPPORT)==0x40) ;
  COMPILE_TIME_ASSERT(sizeof(MMADDRESS_NODE)==0x14) ;
  COMPILE_TIME_ASSERT(sizeof(MM_AVL_TABLE)==0x20) ;

  COMPILE_TIME_ASSERT(sizeof(KAPC_STATE)==0x18) ;
  COMPILE_TIME_ASSERT(sizeof(KAPC)==0x30) ;
  COMPILE_TIME_ASSERT(sizeof(KTRAP_FRAME)==0x8C) ;
  COMPILE_TIME_ASSERT(sizeof(KTHREAD_WXP)==0x1C0) ;
  COMPILE_TIME_ASSERT(sizeof(ETHREAD_WXP)==0x258) ;

  COMPILE_TIME_ASSERT(sizeof(KPROCESS_WXP)==0x06C) ;
  COMPILE_TIME_ASSERT(sizeof(KPROCESS_VISTA)==0x080) ;

  COMPILE_TIME_ASSERT(sizeof(EPROCESS_WXP)==0x260) ;
  COMPILE_TIME_ASSERT(sizeof(EPROCESS_VISTA)==0x270) ;

  SymFile_CloseFile (hSymFile) ;

  return STATUS_SUCCESS ;
}
